{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f47909f",
   "metadata": {},
   "source": [
    "# 06 Distributed Computing\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/sharded-computation.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "45ed1b71",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c3ed583",
   "metadata": {},
   "source": [
    "SPMD is a parallelism technique where the same computation (e.g., a forward pass of a neural network) can be run on different input data (e.g., batch inputs) in parallel on different devices.\n",
    "\n",
    "JAX has three SPMD strategies:\n",
    "\n",
    "- Automatic sharding in JIT contexts.\n",
    "- Explicit sharding which is similar to automatic sharding, but makes the sharding of an array part of the JAX type system as something queryable available at trace-time.\n",
    "  - The compiler still decides how to apply SPMD, but is constrained by user-supplied shardings.\n",
    "- Fully manual sharding with per-device code and explicit communication collectives\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f1391413",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\n"
     ]
    }
   ],
   "source": [
    "# Tell JAX to create 8 logical CPUs, corresponding to 8 OS threads.\n",
    "jax.config.update(\"jax_num_cpu_devices\", 8)\n",
    "print(jax.devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7314ac5",
   "metadata": {},
   "source": [
    "## Data sharding\n",
    "\n",
    "`jax.Array` represents an array in physical memory spanning one or more devices. The type is designed with distributed data and computation in mind.\n",
    "\n",
    "Every `jax.Array` has a `jax.sharding.Sharding` object, which describes which shard of the global data is required by each device.\n",
    "\n",
    "In simple cases, arrays are sharded on a single device:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c26fef1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{CpuDevice(id=0)}\n",
      "SingleDeviceSharding(device=CpuDevice(id=0), memory_kind=unpinned_host)\n"
     ]
    }
   ],
   "source": [
    "a = jnp.arange(32.0).reshape(4, 8)\n",
    "print(a.devices())\n",
    "print(a.sharding)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83cc773a",
   "metadata": {},
   "source": [
    "We can use `jax.debug.visualize_array_sharding()` to visualize how `a` is stored in memory of a single device:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ae32e9f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────┐\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                     CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                      │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "└────────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┌────────────────────────────────────────────────┐\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                     CPU \u001b[1;36m0\u001b[0m                      │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "│                                                │\n",
       "└────────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "jax.debug.visualize_array_sharding(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e5c543a",
   "metadata": {},
   "source": [
    "To create an array with non-trivial sharding, we can define an array with a `jax.sharding` specification and pass this to `jax.device_put()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "77dbbb32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────┬──────────┬──────────┬──────────┐\n",
       "│          │          │          │          │\n",
       "│  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>   │  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>   │  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>   │  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>   │\n",
       "│          │          │          │          │\n",
       "│          │          │          │          │\n",
       "├──────────┼──────────┼──────────┼──────────┤\n",
       "│          │          │          │          │\n",
       "│  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>   │  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>   │  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>   │  CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>   │\n",
       "│          │          │          │          │\n",
       "│          │          │          │          │\n",
       "└──────────┴──────────┴──────────┴──────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┌──────────┬──────────┬──────────┬──────────┐\n",
       "│          │          │          │          │\n",
       "│  CPU \u001b[1;36m0\u001b[0m   │  CPU \u001b[1;36m1\u001b[0m   │  CPU \u001b[1;36m2\u001b[0m   │  CPU \u001b[1;36m3\u001b[0m   │\n",
       "│          │          │          │          │\n",
       "│          │          │          │          │\n",
       "├──────────┼──────────┼──────────┼──────────┤\n",
       "│          │          │          │          │\n",
       "│  CPU \u001b[1;36m4\u001b[0m   │  CPU \u001b[1;36m5\u001b[0m   │  CPU \u001b[1;36m6\u001b[0m   │  CPU \u001b[1;36m7\u001b[0m   │\n",
       "│          │          │          │          │\n",
       "│          │          │          │          │\n",
       "└──────────┴──────────┴──────────┴──────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Make a (2, 4) grid of devices and label X-axis and Y-axis.\n",
    "mesh = jax.make_mesh((2, 4), (\"x\", \"y\"))\n",
    "\n",
    "# Create a specification to shard an array across the mesh.\n",
    "# P(\"x\", \"y\") means split the first dimension of an array across\n",
    "# X-axis and second dimension across Y-axis.\n",
    "sharding = jax.sharding.NamedSharding(mesh, jax.P(\"x\", \"y\"))\n",
    "\n",
    "# Shard the array across the mesh. It will be split into 2 chunks\n",
    "# across the X-axis and 4 chunks across the Y-axis.\n",
    "a = jnp.arange(32.0).reshape(4, 8)\n",
    "a_sharded = jax.device_put(a, sharding)\n",
    "\n",
    "jax.debug.visualize_array_sharding(a_sharded)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b74f7653",
   "metadata": {},
   "source": [
    "## Automatic parallelism in JIT\n",
    "\n",
    "Once we have shared data, the easiest way to do parallel computation is to simply pass the data to a JIT compiled function.\n",
    "\n",
    "The XLA compiler includes heuristics for optimizing computations across multiple devices. We just need to specify how our input and output data is sharded.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "01d34855",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shardings match: True\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f_elementwise(x):\n",
    "    \"\"\"Apply 2 * sin(x) + 1 to an array in element-wise fashion\"\"\"\n",
    "    return 2 * jnp.sin(x) + 1\n",
    "\n",
    "\n",
    "result = f_elementwise(a_sharded)\n",
    "print(\"Shardings match:\", a_sharded.sharding == result.sharding)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9c10075",
   "metadata": {},
   "source": [
    "Here is another example:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "55159f1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┐\n",
       "│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
       "└───────┴───────┴───────┴───────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┌───────┬───────┬───────┬───────┐\n",
       "│CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m│CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m│CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m│CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
       "└───────┴───────┴───────┴───────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f_contract(x):\n",
    "    return jnp.sum(x, axis=0)\n",
    "\n",
    "\n",
    "result = f_contract(a_sharded)\n",
    "jax.debug.visualize_array_sharding(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31533229",
   "metadata": {},
   "source": [
    "The first two elements are replicated on devices 0 and 4, the second two elements on devices 1 and 5, and so on.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dae706fb",
   "metadata": {},
   "source": [
    "## Explicit sharding\n",
    "\n",
    "The key idea is that the JAX-level type of a value includes a description how the data is sharded.\n",
    "\n",
    "The JAX-level type is effectively the information we have access to in a JIT function.\n",
    "\n",
    "We can query the JAX-level type of any JAX value, Numpy array, or Python scalar with `jax.typeof()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1b26d2a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JAX-level type: int32[8]\n",
      "JAX-level type: int32[8]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([ 0,  2,  4,  6,  8, 10, 12, 14], dtype=int32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = jnp.arange(8)\n",
    "\n",
    "# Use jax.typeof() to query the JAX-level type\n",
    "print(f\"JAX-level type:\", jax.typeof(a))\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def f(a):\n",
    "    # jax.typeof() also works in JIT contexts\n",
    "    print(f\"JAX-level type:\", jax.typeof(a))\n",
    "    return 2 * a\n",
    "\n",
    "\n",
    "f(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d4c09c6",
   "metadata": {},
   "source": [
    "To see sharding in the type:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87aff2a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a type: int32[4,2]\n",
      "a_sharded type: int32[4@X,2]\n"
     ]
    }
   ],
   "source": [
    "from jax.sharding import AxisType\n",
    "\n",
    "mesh = jax.make_mesh((2, 4), (\"X\", \"Y\"), axis_types=(AxisType.Explicit, AxisType.Explicit))\n",
    "\n",
    "a = jnp.arange(8).reshape(4, 2)\n",
    "\n",
    "# P(\"X\", None) shards axis 0 across the X mesh (2 shards) and replicates each shard\n",
    "# across the Y mesh (4 replicas per shard). Y replicas hold identical data, so any\n",
    "# computation that doesn't differentiate across Y will be redundant across those devices.\n",
    "a_sharded = jax.device_put(a, jax.NamedSharding(mesh, jax.P(\"X\", None)))\n",
    "\n",
    "print(f\"a type: {jax.typeof(a)}\")\n",
    "print(f\"a_sharded type: {jax.typeof(a_sharded)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bae17fd",
   "metadata": {},
   "source": [
    "`int32[4@X,2]` means a 4 by 2 array of int32s, where the first dimension is sharded across the X-axis of the mesh. The array is replicated across the other mesh axes.\n",
    "\n",
    "In this case, each device will receive a 2 by 2 array. The computation will be replicated across the Y-axis of the mesh.\n",
    "\n",
    "This is an example of adding two sharded arrays:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f024bd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a sharding: int32[4@X,1]\n",
      "b sharding: int32[1,8@Y]\n",
      "ans sharding: int32[4@X,8@Y]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([[ 0,  1,  2,  3,  4,  5,  6,  7],\n",
       "       [ 1,  2,  3,  4,  5,  6,  7,  8],\n",
       "       [ 2,  3,  4,  5,  6,  7,  8,  9],\n",
       "       [ 3,  4,  5,  6,  7,  8,  9, 10]], dtype=int32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Explicitly shard along X-axis\n",
    "a0 = jax.device_put(jnp.arange(4).reshape(4, 1), jax.NamedSharding(mesh, jax.P(\"X\", None)))\n",
    "\n",
    "# Explicitly shard along Y-axis\n",
    "a1 = jax.device_put(jnp.arange(8).reshape(1, 8), jax.NamedSharding(mesh, jax.P(None, \"Y\")))\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def add(a, b):\n",
    "    \"\"\"Add two arrays. Will use Numpy style broadcasting.\"\"\"\n",
    "    ans = a + b\n",
    "    print(f\"a sharding: {jax.typeof(a)}\")\n",
    "    print(f\"b sharding: {jax.typeof(b)}\")\n",
    "    print(f\"ans sharding: {jax.typeof(ans)}\")\n",
    "    return ans\n",
    "\n",
    "\n",
    "# Result will be sharded across both X and Y axes.\n",
    "add(a0, a1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd836cb",
   "metadata": {},
   "source": [
    "Shardings propagate deterministically at trace time and we can query them at trace time.\n",
    "\n",
    "## Manual parallelism\n",
    "\n",
    "In automatic parallelism, we write a function as if we are operation on the full dataset and JIT will split the computation across multiple devices.\n",
    "\n",
    "In manual parallelism with `jax.shard_map()`, we write a function that will handle a single shard of data and `shard_map` will construct the full function:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1c879a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.          2.682942    2.818595    1.28224    -0.513605   -0.9178486\n",
      "  0.44116902  2.3139732   2.9787164   1.824237   -0.08804226 -0.99998045\n",
      " -0.07314587  1.840334    2.9812148   2.3005757   0.42419338 -0.92279494\n",
      " -0.50197446  1.2997544   2.8258905   2.6733112   0.98229736 -0.69244087\n",
      " -0.81115675  0.7352965   2.525117    2.912752    1.5418116  -0.32726777\n",
      " -0.97606325  0.19192469]\n"
     ]
    }
   ],
   "source": [
    "mesh = jax.make_mesh((8,), (\"X\",))\n",
    "\n",
    "\n",
    "def f_elementwise(x):\n",
    "    \"\"\"Apply 2 * sin(x) + 1 to an array in element-wise fashion\"\"\"\n",
    "    return 2 * jnp.sin(x) + 1\n",
    "\n",
    "\n",
    "a = jnp.arange(32)\n",
    "\n",
    "# jax.shard_map() will take a per-shard function and create the full function.\n",
    "f_elementwise_sharded = jax.shard_map(\n",
    "    f_elementwise,\n",
    "    mesh=mesh,\n",
    "    in_specs=jax.P(\"X\"),  # Specify what axis to shard the input array along\n",
    "    out_specs=jax.P(\"X\"),  # Specify what axis to shard the returned output array along\n",
    ")\n",
    "\n",
    "af_sharded = f_elementwise_sharded(a)\n",
    "print(af_sharded)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28de4c53",
   "metadata": {},
   "source": [
    "For reference, `jax.shard_map()` can work inside a JIT function if needed.\n",
    "\n",
    "The function only sees a single portion of the data:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "47180b97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Global shape: x.shape=(32,)\n",
      "Device-local shape: x.shape=(4,)\n"
     ]
    }
   ],
   "source": [
    "x = jnp.arange(32)\n",
    "print(f\"Global shape: {x.shape=}\")\n",
    "\n",
    "\n",
    "def f(x):\n",
    "    print(f\"Device-local shape: {x.shape=}\")\n",
    "    return x * 2\n",
    "\n",
    "\n",
    "y = jax.shard_map(f, mesh=mesh, in_specs=jax.P(\"X\"), out_specs=jax.P(\"X\"))(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a95ce8a9",
   "metadata": {},
   "source": [
    "Because the functions only see the device-local portion of the data, it means that aggregation-like functions require some extra thought.\n",
    "\n",
    "For example, consider summation:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "77bf48f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  6  22  38  54  70  86 102 118]\n"
     ]
    }
   ],
   "source": [
    "x = jnp.arange(32)\n",
    "\n",
    "\n",
    "def f(x):\n",
    "    # keepdims=True will return a 1D array instead of a scalar\n",
    "    return jnp.sum(x, keepdims=True)\n",
    "\n",
    "\n",
    "result = jax.shard_map(f, mesh=mesh, in_specs=jax.P(\"X\"), out_specs=jax.P(\"X\"))(x)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "501aa605",
   "metadata": {},
   "source": [
    "Since `f` operates individually on each shard of the data, we never have any logic to sum across shards. We can explicitly handle this with collective operations like `jax.lax.psum()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "74483c49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "496\n"
     ]
    }
   ],
   "source": [
    "x = jnp.arange(32)\n",
    "\n",
    "\n",
    "def f(x):\n",
    "    shard_sum = x.sum()\n",
    "\n",
    "    # Performs an all-reduce style sum over all shards along X-axis\n",
    "    return jax.lax.psum(shard_sum, \"X\")\n",
    "\n",
    "\n",
    "# Note that out_specs=jax.P(); since the result is a scalar, we do not shard\n",
    "# the result across the X-axis.\n",
    "result = jax.shard_map(f, mesh=mesh, in_specs=jax.P(\"X\"), out_specs=jax.P())(x)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a9fe08b",
   "metadata": {},
   "source": [
    "## Comparing the three approaches\n",
    "\n",
    "Consider a forward pass on a NN layer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "68178f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = jax.random.key(1701)\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def layer(x, W, b):\n",
    "    \"\"\"Applies a dense nn layer with sigmoid activation.\"\"\"\n",
    "    return jax.nn.sigmoid(x @ W + b)\n",
    "\n",
    "\n",
    "key, x_key, W_key, b_key = jax.random.split(key, 4)\n",
    "\n",
    "x = jax.random.normal(x_key, (32,))\n",
    "W = jax.random.normal(W_key, (32, 4))\n",
    "b = jax.random.normal(b_key, (4,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a8651ad",
   "metadata": {},
   "source": [
    "We can automatically perform this in distributed manner using `jax.jit()` and passing appropriately sharded data:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0363fe13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.01422223 0.9701333  0.00124826 0.08572862]\n"
     ]
    }
   ],
   "source": [
    "mesh = jax.make_mesh((8,), (\"X\",))\n",
    "x_sharded = jax.device_put(x, jax.NamedSharding(mesh, jax.P(\"X\")))\n",
    "W_replicated = jax.device_put(W, jax.NamedSharding(mesh, jax.P()))  # P() = replicate\n",
    "\n",
    "# Every device will see 4 input elements. The weights are replicated across all devices.\n",
    "# After each device completes its forward pass, the outputs will be aggregated together\n",
    "# automatically, since JAX will insert a jax.lax.psum() call.\n",
    "print(layer(x_sharded, W_replicated, b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a811583b",
   "metadata": {},
   "source": [
    "The weights are replicated across each device, and the input is sharded across the devices. The final result is joined together in all-reduce fashion.\n",
    "\n",
    "We can also use explicit sharding:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "20fac159",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x sharding: float32[32@X]\n",
      "W sharding: float32[32,4]\n",
      "b sharding: float32[4]\n",
      "out sharding: float32[4]\n",
      "[0.01422223 0.9701333  0.00124826 0.08572862]\n"
     ]
    }
   ],
   "source": [
    "explicit_mesh = jax.make_mesh((8,), (\"X\",), axis_types=(AxisType.Explicit,))\n",
    "x_sharded = jax.device_put(x, jax.NamedSharding(explicit_mesh, jax.P(\"X\")))\n",
    "W_replicated = jax.device_put(W, jax.NamedSharding(explicit_mesh, jax.P()))\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def layer_auto(x, W, b):\n",
    "    \"\"\"Applies a dense nn layer with sigmoid activation.\"\"\"\n",
    "    print(\"x sharding:\", jax.typeof(x))\n",
    "    print(\"W sharding:\", jax.typeof(W))\n",
    "    print(\"b sharding:\", jax.typeof(b))\n",
    "    out = layer(x, W, b)\n",
    "    print(\"out sharding:\", jax.typeof(out))\n",
    "    return out\n",
    "\n",
    "\n",
    "# Every device will see 4 input elements. The weights are replicated across all devices.\n",
    "# After each device completes its forward pass, the outputs will be aggregated together\n",
    "# automatically, since JAX will insert a jax.lax.psum() call.\n",
    "print(layer_auto(x_sharded, W_replicated, b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87de9e92",
   "metadata": {},
   "source": [
    "The key difference here is that the sharding computation is still automatically handled by the compiler, but the sharding is now part of the JAX-level type system.\n",
    "\n",
    "Finally, we can do the same thing with manual parallelism (i.e., `jax.shard_map()`):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c210a262",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.01422223 0.9701333  0.00124826 0.08572862]\n"
     ]
    }
   ],
   "source": [
    "from functools import partial\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "@partial(\n",
    "    jax.shard_map,\n",
    "    mesh=mesh,\n",
    "    in_specs=(\n",
    "        jax.P(\"X\"),  # Shard `x` along X-axis\n",
    "        jax.P(\"X\", None),  # Shard the rows of `W` along X-axis, columns are replicated\n",
    "        jax.P(None),  # `b` is replicated across all devices\n",
    "    ),\n",
    "    out_specs=jax.P(None),  # Output is replicated across all devices\n",
    ")\n",
    "def layer_sharded(x, W, b):\n",
    "    # Each device computes its local partial (x_i @ W_i) in R^M.\n",
    "    # psum(..., \"X\") then does an elementwise all-reduce (sum) across devices,\n",
    "    # yielding the full x @ W on every device. Adding b and sigmoid are local.\n",
    "    return jax.nn.sigmoid(jax.lax.psum(x @ W, \"X\") + b)\n",
    "\n",
    "\n",
    "print(layer_sharded(x, W, b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1385c866",
   "metadata": {},
   "source": [
    "## Distributed Training\n",
    "\n",
    "We can shard model weights across device and incrementally update them during training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "513008dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = jax.random.key(1701)\n",
    "mesh = jax.make_mesh((8,), (\"X\",))\n",
    "\n",
    "B = 128 * jax.device_count()\n",
    "D = 32\n",
    "LEARNING_RATE = 0.01\n",
    "\n",
    "# Generate dataset\n",
    "X = jax.random.normal(key, (B, D))\n",
    "w_true = jax.random.normal(key, (D, 1))  # True weights for label generation\n",
    "y = (X @ w_true > 0).astype(jnp.float32)  # Labels: 1 if positive, 0 else\n",
    "\n",
    "# Shard data and parameters across X-axis\n",
    "data_sharding = jax.NamedSharding(mesh, jax.P(\"X\", None))\n",
    "params_sharding = jax.NamedSharding(mesh, jax.P())\n",
    "\n",
    "X_sharded = jax.device_put(X, data_sharding)\n",
    "y_sharded = jax.device_put(y, data_sharding)\n",
    "\n",
    "params = {\"w\": jnp.zeros((D, 1)), \"b\": jnp.zeros((1,))}\n",
    "params_sharded = jax.tree.map(lambda a: jax.device_put(a, params_sharding), params)\n",
    "\n",
    "\n",
    "@partial(\n",
    "    jax.jit,\n",
    "    in_shardings=(params_sharding, data_sharding, data_sharding),\n",
    "    out_shardings=params_sharding,\n",
    ")\n",
    "def train_step(p, x, y):\n",
    "    \"\"\"\n",
    "    Perform one step of training with data-parallel sharding.\n",
    "\n",
    "    Args:\n",
    "    - p: Model parameters replicated across all devices.\n",
    "    - x: Input features of the batch sharded along the batch dimension (X-axis).\n",
    "         In this example, each device recieves 128 examples per step (128 * 8 / 8).\n",
    "         Since there are 32 features, each device effectively has a 128x8 matrix.\n",
    "    - y: Target labels for the batch.\n",
    "    \"\"\"\n",
    "\n",
    "    def loss_fn(p):\n",
    "        logits = x @ p[\"w\"] + p[\"b\"]\n",
    "        preds = jax.nn.sigmoid(logits)\n",
    "        return jnp.mean((preds - y) ** 2)\n",
    "\n",
    "    # Update weights & bias with loss gradient direction\n",
    "    grad = jax.grad(loss_fn)(p)\n",
    "    return jax.tree.map(lambda a, g: a - LEARNING_RATE * g, p, grad)\n",
    "\n",
    "\n",
    "# Train for 100 steps\n",
    "for _ in range(100):\n",
    "    params = train_step(params_sharded, X_sharded, y_sharded)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3b05306",
   "metadata": {},
   "source": [
    "The model’s weights are replicated on every device, so each device holds a full copy and updates it in sync.\n",
    "\n",
    "The input batch is sharded along the X-axis, so each device processes a different slice of examples, computes local gradients, and JAX automatically all-reduces those gradients so every device applies the same weight update.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
